/**
 * @file MinRes.h
 * @author LSJ (you@domain.com)
 * @brief 
 * @version 0.1
 * @date 2021-01-12
 * 
 * @copyright Copyright (c) 2021
 * 
 */
#ifndef _MinRes_H_
#define _MinRes_H_
#include <Eigen/Sparse>
#include <Eigen/IterativeLinearSolvers>
namespace Eigen
{

namespace internal {
template<typename MatrixType, typename Rhs, typename Dest, typename Preconditioner>
void minres(const MatrixType& mat, const Rhs& rhs, Dest& x,
              const Preconditioner& precond, Index& iters,
              typename Dest::RealScalar& tol_error)
{
    using std::sqrt;
    typedef typename Dest::RealScalar RealScalar;
    typedef typename Dest::Scalar Scalar;
    typedef Matrix<Scalar,Dynamic,1> VectorType;

    // Check for zero rhs
    const RealScalar rhsNorm2(rhs.squaredNorm());
    if(rhsNorm2 == 0)
    {
        x.setZero();
        iters = 0;
        tol_error = 0;
        return;
    }
            
    // initialize
    const Index maxIters(iters);  // initialize maxIters to iters
    const Index N(mat.cols());    // the size of the matrix
    const RealScalar threshold2(tol_error*tol_error*rhsNorm2); // convergence threshold (compared to residualNorm2)
            
    // Initialize preconditioned Lanczos
    VectorType v_old(N); // will be initialized inside loop
    VectorType v( VectorType::Zero(N) ); //initialize v
    VectorType v_new(rhs-mat*x); //initialize v_new
    RealScalar residualNorm2(v_new.squaredNorm());
    VectorType w(N); // will be initialized inside loop
    VectorType w_new(precond.solve(v_new)); // initialize w_new
    // will be initialized inside loop
    RealScalar beta_new2(v_new.dot(w_new));
    eigen_assert(beta_new2 >= 0.0 && "PRECONDITIONER IS NOT POSITIVE DEFINITE");
    RealScalar beta_new(sqrt(beta_new2));
    RealScalar beta_one(beta_new);
    v_new /= beta_new;
    w_new /= beta_new;
    // Initialize other variables
    RealScalar c(1.0); // the cosine of the Givens rotation
    RealScalar c_old(1.0);
    RealScalar s(0.0); // the sine of the Givens rotation
    RealScalar s_old(0.0); // the sine of the Givens rotation
    VectorType p_oold(N); // will be initialized in loop
    VectorType p_old(VectorType::Zero(N)); // initialize p_old=0
    VectorType p(p_old); // initialize p=0
    RealScalar eta(1.0);
                        
    iters = 0; // reset iters
    while ( iters < maxIters )
            {
                const RealScalar beta(beta_new);
                v_old = v; // update: at first time step, this makes v_old = 0 so value of beta doesn't matter
//                const VectorType v_old(v); // NOT SURE IF CREATING v_old EVERY ITERATION IS EFFICIENT
                v = v_new; // update
                w = w_new; // update
//                const VectorType w(w_new); // NOT SURE IF CREATING w EVERY ITERATION IS EFFICIENT
                v_new.noalias() = mat*w - beta*v_old; // compute v_new
                const RealScalar alpha = v_new.dot(w);
                v_new -= alpha*v; // overwrite v_new
                w_new = precond.solve(v_new); // overwrite w_new
                beta_new2 = v_new.dot(w_new); // compute beta_new
                eigen_assert(beta_new2 >= 0.0 && "PRECONDITIONER IS NOT POSITIVE DEFINITE");
                beta_new = sqrt(beta_new2); // compute beta_new
                v_new /= beta_new; // overwrite v_new for next iteration
                w_new /= beta_new; // overwrite w_new for next iteration
                
                // Givens rotation
                const RealScalar r2 =s*alpha+c*c_old*beta; // s, s_old, c and c_old are still from previous iteration
                const RealScalar r3 =s_old*beta; // s, s_old, c and c_old are still from previous iteration
                const RealScalar r1_hat=c*alpha-c_old*s*beta;
                const RealScalar r1 =sqrt( std::pow(r1_hat,2) + std::pow(beta_new,2) );
                c_old = c; // store for next iteration
                s_old = s; // store for next iteration
                c=r1_hat/r1; // new cosine
                s=beta_new/r1; // new sine
                
                // Update solution
                p_oold = p_old;
//                const VectorType p_oold(p_old); // NOT SURE IF CREATING p_oold EVERY ITERATION IS EFFICIENT
                p_old = p;
                p.noalias()=(w-r2*p_old-r3*p_oold) /r1; // IS NOALIAS REQUIRED?
                x += beta_one*c*eta*p;
                
                /* Update the squared residual. Note that this is the estimated residual.
                The real residual |Ax-b|^2 may be slightly larger */
                residualNorm2 *= s*s;
                
                if ( residualNorm2 < threshold2)
                {
                    break;
                }
                
                eta=-s*eta; // update eta
                iters++; // increment iteration number (for output purposes)
            }
                   
    /* Compute error. Note that this is the estimated error. The real 
    error |Ax-b|/|b| may be slightly larger */
    tol_error = std::sqrt(residualNorm2 / rhsNorm2);
}
}

template< typename _MatrixType, typename _Preconditioner = IdentityPreconditioner,int _UpLo = 1>
class MinRes;

namespace internal
{
template< typename _MatrixType, typename _Preconditioner,int _UpLo>
struct traits<MinRes<_MatrixType,_Preconditioner,_UpLo> >
{
  typedef _MatrixType MatrixType;
  typedef _Preconditioner Preconditioner;
};    
}


 /** \ingroup IterativeLinearSolvers_Module
     * \brief A minimal residual solver for sparse symmetric problems
     *
     * This class allows to solve for A.x = b sparse linear problems using the MINRES algorithm
     * of Paige and Saunders (1975). The sparse matrix A must be symmetric (possibly indefinite).
     * The vectors x and b can be either dense or sparse.
     *
     * \tparam _MatrixType the type of the sparse matrix A, can be a dense or a sparse matrix.
     * \tparam _Preconditioner the type of the preconditioner. Default is DiagonalPreconditioner
*/
template< typename _MatrixType, typename _Preconditioner,int _UpLo>
class MinRes : public IterativeSolverBase<MinRes<_MatrixType,_Preconditioner,_UpLo> >
{
    typedef IterativeSolverBase<MinRes> Base;
    using Base::matrix;
    using Base::m_error;
    using Base::m_iterations;
    using Base::m_info;
    using Base::m_isInitialized;
    public:
    typedef _MatrixType MatrixType;
    typedef typename MatrixType::Scalar Scalar;
    typedef typename MatrixType::RealScalar RealScalar;
    typedef _Preconditioner Preconditioner;
    public:
    MinRes():Base(){};
    template<typename MatrixDerived>
    explicit MinRes(const EigenBase<MatrixDerived>& A) : Base(A.derived()) {}
    MinRes(const Preconditioner& preconditioner){_preconditioner = preconditioner;};
    ~MinRes(){};
    enum {UpLo = _UpLo};
    private:
    Preconditioner _preconditioner;
    public:
    template<typename Rhs,typename Dest>
    void _solve_with_guess_impl(const Rhs& b, Dest& x) const
    {    
        typedef typename Base::MatrixWrapper MatrixWrapper;
            typedef typename Base::ActualMatrixType ActualMatrixType;
            enum {
              TransposeInput  =   (!MatrixWrapper::MatrixFree)
                              &&  (UpLo==(Lower|Upper))
                              &&  (!MatrixType::IsRowMajor)
                              &&  (!NumTraits<Scalar>::IsComplex)
            };
            typedef typename internal::conditional<TransposeInput,Transpose<const ActualMatrixType>, ActualMatrixType const&>::type RowMajorWrapper;
            EIGEN_STATIC_ASSERT(EIGEN_IMPLIES(MatrixWrapper::MatrixFree,UpLo==(Lower|Upper)),MATRIX_FREE_CONJUGATE_GRADIENT_IS_COMPATIBLE_WITH_UPPER_UNION_LOWER_MODE_ONLY);
            typedef typename internal::conditional<UpLo==(Lower|Upper),
                                                  RowMajorWrapper,
                                                  typename MatrixWrapper::template ConstSelfAdjointViewReturnType<UpLo>::Type
                                            >::type SelfAdjointWrapper;
        for(Index j=0; j<b.cols(); ++j)
        {
        m_iterations = Base::maxIterations();
        m_error = Base::m_tolerance;
        typename Dest::ColXpr xj(x,j);
        internal::minres(SelfAdjointWrapper(matrix()), b.col(j), xj,
                                _preconditioner, m_iterations, m_error);
        }
        m_isInitialized = true;
        m_info = m_error <= Base::m_tolerance ? Success : NoConvergence;
    };

    using Base::_solve_impl;
    template<typename Rhs,typename Dest>
    void _solve_impl(const MatrixBase<Rhs>& b, Dest& x) const
    {
        x.resize(this->rows(),b.cols());
        x.setZero();
        _solve_with_guess_impl(b,x);
     };
};
}
/** \ingroup The preconditioner for Stokes system
     * \brief The preconditioner for Stokes system
     *
     * This class is used to preconditioning the discrete sparse or dense system from stokes eqution. 
     * The constructor of this class need two parameters,A and Q. A can be selected by the matrix discrete from
     * poission equation and one of the choices of Q is the mass matrix,or shur-competetion matrix.
     *
     * \tparam _Scalar the type of the scalar.
     * \tparam _MatrixType the type of the sparse matrix , can be a dense or a sparse matrix.
*/
template< typename _Scalar,typename _MatrixType>
class StokesPreconditioner
{
    public:
    typedef _MatrixType MatrixType;
    typedef _Scalar Scalar;
    typedef Eigen::Matrix<Scalar,Eigen::Dynamic,1> VectorType;
    typedef typename VectorType::StorageIndex StorageIndex;
    enum {
      ColsAtCompileTime = Eigen::Dynamic,
      MaxColsAtCompileTime = Eigen::Dynamic
    };
    public:
    StokesPreconditioner(){};
    StokesPreconditioner(const MatrixType& A,const MatrixType& Q) : _A(A) ,_Q(Q){};

    ~StokesPreconditioner(){};
    
    StokesPreconditioner& operator=(const StokesPreconditioner& _other)
    {
        this->_A = _other._A;
        this->_Q = _other._Q;
        return *this;
    }
    Eigen::Index cols() const { return 2 *  _A.cols() + _Q.cols(); }
    template<typename Rhs, typename Dest>
    void _solve_impl(const Rhs& b, Dest& x) const
    {
        typedef typename Dest::RealScalar RealScalar;
        typedef typename Dest::Scalar Scalar;
        typedef Eigen::ConjugateGradient<MatrixType,Eigen::Lower|Eigen::Upper,Eigen::IncompleteCholesky<Scalar> > CGSolver;
        typedef Eigen::Matrix<Scalar,Eigen::Dynamic,1> VectorType;
        Eigen::Index n_A = _A.cols();
        Eigen::Index n_Q = _Q.cols();
        VectorType vx = b.segment(0,n_A);
        VectorType vy = b.segment(n_A,n_A);
        VectorType vp = b.segment(2 * n_A,n_Q);
        CGSolver cg1;
        CGSolver cg2;
        cg1.compute(_A);
        cg2.compute(_Q);
        VectorType xx = cg1.solve(vx);
        VectorType xy = cg1.solve(vy);
        VectorType xp = cg2.solve(vp);

        x.segment(0,n_A) = xx;
        x.segment(n_A,n_A) = xy;
        x.segment(2*n_A,n_Q) = xp;
    };
    Eigen::ComputationInfo info() { return Eigen::Success; }

    template<typename MatrixType>
    StokesPreconditioner& compute(const MatrixType& ) { return *this; }

    template<typename Rhs> inline const Eigen::Solve<StokesPreconditioner, Rhs>
    solve(const Eigen::MatrixBase<Rhs>& b) const
    {
      eigen_assert(_A.cols() * 2 + _Q.cols() ==b.rows()
                && "DiagonalPreconditioner::solve(): invalid number of rows of the right hand side matrix b");
      return Eigen::Solve<StokesPreconditioner, Rhs>(*this, b.derived());
    }
    MatrixType _A;
    MatrixType _Q;
};
template< typename _Scalar,typename _MatrixType = Eigen::SparseMatrix<_Scalar> >
class StokesPreconditioner;

#else
#endif 